package main

import (
	"context"
	"encoding/json"
	"fmt"
	"happy-ai/config"
	"log"
	"os"

	// Eino & Ark 扩展
	"github.com/cloudwego/eino-ext/components/model/ark"
	"github.com/cloudwego/eino/schema"

	// MCP 客户端库
	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/mcp"
)

func main() {
	err := config.Init("")
	if err != nil {
		return
	}
	ctx := context.Background()
	log.Println("启动 Eino Agent (Backend: Volcano Ark)...")

	// ======================================================================================
	// 1. 初始化 MCP Clients
	// ======================================================================================

	initRequest := mcp.InitializeRequest{
		Request: mcp.Request{
			Method: string(mcp.MethodInitialize),
		},
		Params: mcp.InitializeParams{
			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
			ClientInfo: mcp.Implementation{
				Name:    "eino-agent-client",
				Version: "1.0.0",
			},
			Capabilities: mcp.ClientCapabilities{},
		},
	}

	// --- Sales Server Client ---
	salesClient, err := client.NewStreamableHttpClient(
		"http://localhost:2508/query_sales_data",
	)
	if err != nil {
		log.Fatalf("Failed to create sales mcp client: %v", err)
	}
	defer salesClient.Close()

	if err := salesClient.Start(ctx); err != nil {
		log.Fatalf("Failed to start sales mcp: %v", err)
	}

	// 接收 Initialize 的两个返回值
	salesInitResult, err := salesClient.Initialize(ctx, initRequest)
	if err != nil {
		log.Fatalf("sales client init error: %v", err)
	}
	log.Printf("Sales Server Connected: %s", salesInitResult.ServerInfo.Name)
	// --- AMap Server Client ---
	amapClient, err := client.NewStreamableHttpClient(
		"http://localhost:2507/plan_ebike_route",
	)
	if err != nil {
		log.Fatalf("Failed to create amap mcp client: %v", err)
	}
	defer amapClient.Close()

	if err := amapClient.Start(ctx); err != nil {
		log.Fatalf("Failed to start amap mcp: %v", err)
	}

	amapInitResult, err := amapClient.Initialize(ctx, initRequest)
	if err != nil {
		log.Fatalf("amap client init error: %v", err)
	}
	log.Printf("AMap Server Connected: %s", amapInitResult.ServerInfo.Name)

	// ======================================================================================
	// 2. 定义 Eino 工具 Schema
	// ======================================================================================

	salesParams := map[string]*schema.ParameterInfo{
		"city":    {Type: schema.String, Desc: "城市名称"},
		"sort_by": {Type: schema.String, Desc: "排序字段: sale_total/sale_date", Enum: []string{"sale_total", "sale_date"}},
		"order":   {Type: schema.String, Desc: "排序方式: ASC/DESC", Enum: []string{"ASC", "DESC"}},
		"limit":   {Type: schema.Integer, Desc: "限制数量"},
	}

	salesToolInfo := &schema.ToolInfo{
		Name:        "query_sales_data",
		Desc:        "查询药品销售明细。",
		ParamsOneOf: schema.NewParamsOneOfByParams(salesParams),
	}

	amapParams := map[string]*schema.ParameterInfo{
		"origin":      {Type: schema.String, Desc: "起点坐标", Required: true},
		"destination": {Type: schema.String, Desc: "终点坐标", Required: true},
		"waypoints":   {Type: schema.String, Desc: "途径点", Required: true},
	}

	amapToolInfo := &schema.ToolInfo{
		Name:        "plan_ebike_route",
		Desc:        "规划骑行路径。",
		ParamsOneOf: schema.NewParamsOneOfByParams(amapParams),
	}

	// ======================================================================================
	// 3. 定义工具具体实现 Handler (含修复逻辑)
	// ======================================================================================

	// 辅助函数：从 MCP 响应中提取文本内容
	extractContent := func(content []mcp.Content) string {
		if len(content) == 0 {
			return "No content returned"
		}
		// 尝试断言为 TextContent
		if textContent, ok := content[0].(mcp.TextContent); ok {
			return textContent.Text
		}
		// 如果不是文本（如图片），则序列化为 JSON 字符串作为兜底
		rawBytes, _ := json.Marshal(content)
		return string(rawBytes)
	}

	callSalesTool := func(argsJSON string) (string, error) {
		var args struct {
			City   string `json:"city"`
			SortBy string `json:"sort_by"`
			Order  string `json:"order"`
			Limit  int    `json:"limit"`
		}
		args.Limit = 20
		args.City = "西安市"
		args.SortBy = "sale_total"
		args.Order = "DESC"

		if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
			return "", fmt.Errorf("args parse error: %v", err)
		}

		req := mcp.CallToolRequest{
			Request: mcp.Request{
				Method: string(mcp.MethodToolsCall),
			},
			Params: mcp.CallToolParams{
				Name: "query_sales_data",
				Arguments: map[string]interface{}{
					"city": args.City, "sort_by": args.SortBy, "order": args.Order, "limit": args.Limit,
				},
			},
		}

		res, err := salesClient.CallTool(ctx, req)
		if err != nil {
			return "", err
		}
		if res.IsError {
			return "", fmt.Errorf("mcp error: %v", res)
		}

		// 修复点：使用辅助函数提取内容
		return extractContent(res.Content), nil
	}

	callAmapTool := func(argsJSON string) (string, error) {
		var args struct {
			Origin      string `json:"origin"`
			Destination string `json:"destination"`
			Waypoints   string `json:"waypoints"`
		}
		if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
			return "", fmt.Errorf("args parse error: %v", err)
		}

		req := mcp.CallToolRequest{
			Request: mcp.Request{
				Method: string(mcp.MethodToolsCall),
			},
			Params: mcp.CallToolParams{
				Name: "plan_ebike_route",
				Arguments: map[string]interface{}{
					"origin": args.Origin, "destination": args.Destination,
					"waypoints": args.Waypoints,
				},
			},
		}

		res, err := amapClient.CallTool(ctx, req)
		if err != nil {
			return "", err
		}
		if res.IsError {
			return "", fmt.Errorf("mcp error: %v", res)
		}

		// 修复点：使用辅助函数提取内容
		return extractContent(res.Content), nil
	}

	// ======================================================================================
	// 4. 初始化火山引擎 Ark Model
	// ======================================================================================
	arkApiKey := os.Getenv("ARK_API_KEY")
	arkModelID := os.Getenv("ARK_MODEL_ID")
	if arkApiKey == "" || arkModelID == "" {
		arkApiKey = "d389cada-93cd-4b23-930d-7b187b4b1cff"
		arkModelID = "doubao-seed-1-6-251015"
	}

	chatModel, err := ark.NewChatModel(ctx, &ark.ChatModelConfig{
		APIKey: arkApiKey,
		Model:  arkModelID,
		//AccessKey: "AKLTNmZiNjQ4MDc5OTI5NDY3Njk0YzljMDhkOGYwZDZmZjg",
		//SecretKey: "TWpVM01tVTBOMlpsWWpkak5ERTFOamswTVdZeE5XTmlNRGRoWmpBME5qTQ==",
	})
	if err != nil {
		log.Fatalf("Ark model init failed: %v", err)
	}

	err = chatModel.BindTools([]*schema.ToolInfo{salesToolInfo, amapToolInfo})
	if err != nil {
		log.Fatalf("BindTools failed: %v", err)
	}

	// ======================================================================================
	// 5. 执行 Agent ReAct Loop
	// ======================================================================================

	systemPrompt := `你是一个专业的医药销售代表助理。
你的核心任务是帮助销售代表优化拜访路线。

当前场景：
1. 请查询 "西安市" 销量最低 (sale_total ASC) 的前 3 家药店。
2. 假设销售代表当前位置在：108.93977,34.341574 (未央区政府附近)。
3. 请为这 3 家药店规划骑行路线，并计算总耗时。

请一步步思考：先调用数据查询工具，获取药店的 location 坐标，然后针对所有药店调用路径规划工具, 终点选择距离当前位置最远的药店, 其余作为途径点(waypoints)。`

	messages := []*schema.Message{
		{Role: schema.System, Content: systemPrompt},
		{Role: schema.User, Content: `
#帮我制作下一个工作日的巡店攻略，考虑出行时间和路线，以及天气状况路线规划。
#制作网页，页面包含天气卡片展示3天的天气详情以及巡店小贴士，展示旅行计划， 安排3家药店。
#根据已定的起始位置, 制作网页地图自定义绘制巡店路线, 出行方案选择骑行。
##网页使用简约美观页面风格，并且要呈现分享二维码, 以供在手机APP中使用导航。
#行程规划结果在高德地图app展示，并集成到h5页面中。
#将最终生成文件的html源码发出来。
`},
	}

	const maxRounds = 10
	log.Println("=== 开始对话循环 ===")

	for i := 0; i < maxRounds; i++ {
		msg, err := chatModel.Generate(ctx, messages)
		if err != nil {
			log.Fatalf("Generate error: %v", err)
		}

		log.Printf("\n[Round %d] Role: %s\nContent: %s\n", i+1, msg.Role, msg.Content)

		messages = append(messages, msg)

		if len(msg.ToolCalls) == 0 {
			log.Println("任务完成或无工具调用。")
			break
		}

		for _, call := range msg.ToolCalls {
			log.Printf(">>> Calling Tool: %s, Args: %s", call.Function.Name, call.Function.Arguments)

			var result string
			var toolErr error

			switch call.Function.Name {
			case "query_sales_data":
				result, toolErr = callSalesTool(call.Function.Arguments)
			case "plan_ebike_route":
				result, toolErr = callAmapTool(call.Function.Arguments)
			default:
				result = "Error: Unknown tool"
			}

			if toolErr != nil {
				result = fmt.Sprintf("Tool Execution Error: %v", toolErr)
			}

			logMsg := result
			if len(logMsg) > 200 {
				logMsg = logMsg[:200] + "..."
			}
			log.Printf("<<< Tool Result: %s", logMsg)

			messages = append(messages, &schema.Message{
				Role:       schema.Tool,
				Content:    result,
				ToolCallID: call.ID,
			})
		}
	}
}
